#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@Author: yinwh
@Time: 2025-06-08 10:25
@File: client.py
@Version: 1.0.0
@Description: 
@Copyright: (c) 2025 by yinwh. All rights reserved.
"""
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv
import os
from tools import *
from typing import List, Dict, Any
from utils import function_to_json

_ = load_dotenv(find_dotenv())

SYSTEM_PROMPT = """
你是一个叫不要葱姜蒜的人工智能助手。你的输出应该与用户的语言保持一致。
当用户的问题需要调用工具时，你可以从提供的工具列表中调用适当的工具函数。
"""


class LLMAgent(object):

    def __init__(self, client: OpenAI, model: str = "Qwen/Qwen2.5-32B-Instruct", tools: List = [], verbose: bool = True):
        self.client = client
        self.tools = tools
        self.model = model
        self.messages = [
            {"role": "system", "content": SYSTEM_PROMPT}
        ]
        self.verbose = verbose

    def get_tool_schema(self) -> List[Dict[str, Any]]:
        return [function_to_json(tool) for tool in self.tools]

    def handle_tool_call(self, tool_call):
        function_name = tool_call.function.name
        function_args = tool_call.function.arguments
        function_id = tool_call.id

        function_call_content = eval(f"{function_name}(**{function_args})")

        return {
            "role": "tool",
            "content": function_call_content,
            "tool_call_id": function_id
        }

    def get_completion(self, prompt) -> str:
        self.messages.append({"role": "user", "content": prompt})

        response = self.client.chat.completions.create(
            model = self.model,
            messages = self.messages,
            tools = self.get_tool_schema(),
            stream = False,
        )

        if response.choices[0].message.tool_calls:
            self.messages.append(response.choices[0].message)

            # 处理工具调用
            tool_list = []
            for tool_call in response.choices[0].message.tool_calls:
                # 调用工具并将结果添加到消息列表中。
                tool_response = self.handle_tool_call(tool_call)
                self.messages.append(tool_response)  # 确保工具调用的结果被正确添加
                tool_list.append([tool_call.function.name, tool_call.function.arguments])

            if self.verbose:
                print("调用工具: ", response.choices[0].message.content, tool_list)

            # 重新调用模型以获取最终响应
            response = self.client.chat.completions.create(
                model = self.model,
                messages = self.messages,
                tools = self.get_tool_schema(),
                stream = False,
            )

        # 将模型的完成响应添加到消息列表中
        self.messages.append({
            "role": "assistant",
            "content": response.choices[0].message.content,
        })
        return response.choices[0].message.content



